Add test for autodifferentiating hydrostatic turbulence simulation #3867
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
|
Here's the error that I get right now when running this locally: Farther down the error message says: And the stacktrace says: |
|
It's coming from this annoying function: Oceananigans.jl/src/Fields/field_tuples.jl Lines 56 to 83 in 30e3e05 which we can certainly simplify. We can also play around this avoiding hitting this line altogether, at least for a certain class of models. |
|
Ok problem may be fixed. |
|
oh cool, well in any case this is why we should make PRs/land everything we want to experiment with [and things that fail I can fix] |
100% --- it seems like the problems here are all on the Oceananigans side |
…com/CliMA/Oceananigans.jl into glw/autodiff-hydrostatic-turbulence
|
Is this correct in lines 274=275 of When I run it locally it returns this error: I found a few possible errors locally and can push. |
|
Oceananigans.jl/test/test_enzyme.jl Lines 274 to 278 in 45dbdf8 Yes that's a bug -- fixed now PS to generate a link to code, click the "..." on the left side. You'll get an option to create a permalink. You can edit the link to refer to a range of lines after creating it as I did above (to see what I wrote, attempt to edit my post and you will see the hyperlink)
|
|
Here's another code snipped that I have been using to test this: using Oceananigans
using Oceananigans.Utils: with_tracers
using Random
using Enzyme
Random.seed!(123)
arch = CPU()
Nx = Ny = 32
x = y = (0, 2π)
z = (0, 1)
g = 4^2
c = sqrt(g)
grid = RectilinearGrid(arch, size=(Nx, Ny, 1); x, y, z, topology=(Periodic, Periodic, Bounded))
closure = ScalarDiffusivity(ν=1e-2)
momentum_advection = Centered(order=2)
free_surface = ExplicitFreeSurface(gravitational_acceleration=g)
model = HydrostaticFreeSurfaceModel(; grid, momentum_advection, free_surface, closure)
ϵ(x, y, z) = 2randn() - 1
set!(model, u=ϵ, v=ϵ)
u_init = deepcopy(model.velocities.u)
v_init = deepcopy(model.velocities.v)
Δx = minimum_xspacing(grid)
Δt = 0.01 * Δx / c
for n = 1:10
time_step!(model, Δt; euler=true)
end
u_truth = deepcopy(model.velocities.u)
v_truth = deepcopy(model.velocities.v)
function set_viscosity!(model, viscosity)
new_closure = ScalarDiffusivity(ν=viscosity)
names = ()
new_closure = with_tracers(names, new_closure)
model.closure = new_closure
return nothing
end
function viscous_hydrostatic_turbulence(ν, model, u_init, v_init, Δt, u_truth, v_truth)
# Initialize the model
model.clock.iteration = 0
model.clock.time = 0
#model.clock.last_Δt = Inf
set_viscosity!(model, ν)
#set!(model, u=u_init, v=v_init, η=0)
set!(model, u=u_init, v=v_init)
fill!(parent(model.free_surface.η), 0)
# Step it forward
for n = 1:10
time_step!(model, Δt; euler=true)
end
# Compute the sum square error
u, v, w = model.velocities
Nx, Ny, Nz = size(model.grid)
err = 0.0
for j = 1:Ny, i = 1:Nx
err += @inbounds (u[i, j, 1] - u_truth[i, j, 1])^2 +
(v[i, j, 1] - v_truth[i, j, 1])^2
end
return err::Float64
end
# Use a manual finite difference to compute a gradient
Δν = 1e-6
ν1 = 1.1e-2
ν2 = ν1 + Δν
e1 = viscous_hydrostatic_turbulence(ν1, model, u_init, v_init, Δt, u_truth, v_truth)
e2 = viscous_hydrostatic_turbulence(ν2, model, u_init, v_init, Δt, u_truth, v_truth)
Δe = e2 - e1
ΔeΔν = (e2 - e1) / Δν
@info "Finite difference computed: $ΔeΔν"
@info "Now with autodiff..."
start_time = time_ns()
# Use autodiff to compute a gradient
dmodel = Enzyme.make_zero(model)
dedν = autodiff(set_runtime_activity(Enzyme.Reverse),
viscous_hydrostatic_turbulence,
Active(ν1),
Duplicated(model, dmodel),
Const(u_init),
Const(v_init),
Const(Δt),
Const(u_truth),
Const(v_truth))
@info "Automatically computed: $dedν."
@info "Elapsed time: " * prettytime(1e-9 * (time_ns() - start_time)) |
|
Cool! Two other possible bugs I found: Oceananigans.jl/test/test_enzyme.jl Line 331 in 1522ae4 I think should be Oceananigans.jl/test/test_enzyme.jl Lines 261 to 264 in 1522ae4 produces an out of bounds error with the z axis. I replaced Happy to push those changes. |
I noticed that. Why was it [1][3] before? We can't use |
|
Not exactly sure why it was [1][3], but if |
Hm ok got it, it depends on the status of what's returned |
jlk9
previously approved these changes
Nov 7, 2024
|
Some of the distributed tests are failing, and I wonder if the changes to The stack trace goes through Looking at the function handles: The main difference in the handles is the we see that the Clock object is getting treated as a buffer instead of part of args. The PR doesn't change |
i'll look into this |
I see you have changed Oceananigans.jl/src/DistributedComputations/halo_communication.jl Lines 81 to 85 in 4e22432 In this PR it seems like this extension has been bypassed, and the tests are trying to send halos of field tuples. Maybe it's a useful information to debug |
…egions! is the issue
|
Oceananigans.jl/src/Fields/field_tuples.jl Lines 56 to 72 in 11809b0 Probably not the way to permanently fix it, but I've just added |
glwagner
commented
Nov 10, 2024
glwagner
commented
Nov 12, 2024
Co-authored-by: Gregory L. Wagner <[email protected]>
…glw/autodiff-hydrostatic-turbulence
|
There's currently a circular dependency issue with KernelAbstractions v0.9.30 and some Oceananigans / Enzyme extensions: this seems to be why most GPU tests are failing. |
Its a Julia bug, but are you sure this is why GPU tests fail? |
|
Oh ok. For me running Enzymanigans scripts with KA 9.30 locally produces errors related to compilation. I thought it was due to the circular dependency but maybe it's something else? In any case I've been staying on 9.29. |
I believe you can ignore the |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
This PR adds a test that automatically differentiates a hydrostatic turbulence simulation that uses
Centeredmomentum advection and anExplicitFreeSurface.@wsmoses @jlk9
Previously we had deduced that this model was differentiable, but it looks like there are still some issues. Not sure what the other test was doing. I think we should consider this separately from #3822 since that PR is currently working on the
SplitExplicitFreeSurface.